
"""
红利白马股交易策略实现
"""
import pandas as pd
import numpy as np
from datetime import datetime
from akshare_request import get_stock_zh_a_hist_safe1
import time

# 配置参数
ANALYSIS_CONFIG = {
  'stock_codes':['600900', '601919'],
  'pe_threshold':20, # 市盈率阈值
  'dividend_yield':0.03, # 股息率阈值
  'ma_period':120, # 移动平均周期
  'buy_threshold':0.12, # 买入阈值(相对于MA120的百分比)
  'sell_threshold':0.12, # 卖出阈值(相对于MA120的百分比)
  'start_date':'2015-01-01',
  'end_date':'2025-12-31',
  'enable_backtest':True,
}

class DividendStrategy:
    """
    红利白马股交易策略
    
    参数:
        stock_codes: 股票代码列表，例如 ['600900', '601919']
        pe_threshold: 市盈率阈值，默认20
        dividend_yield: 股息率阈值，默认0.03 (3%)
        ma_period: 移动平均周期，默认120
        buy_threshold: 买入阈值(相对于MA120的百分比)，默认0.12 (12%)
        sell_threshold: 卖出阈值(相对于MA120的百分比)，默认0.12 (12%)
        start_date: 回测开始日期，默认'2015-01-01'
        end_date: 回测结束日期，默认'2025-12-31'
    """
    
    def __init__(self, stock_codes, pe_threshold=20, dividend_yield=0.03, 
                 ma_period=120, buy_threshold=0.12, sell_threshold=0.12,
                 start_date='2015-01-01', end_date='2025-12-31', enable_backtest=True):
        self.stock_codes = stock_codes
        self.pe_threshold = pe_threshold
        self.dividend_yield = dividend_yield
        self.ma_period = ma_period
        self.buy_threshold = buy_threshold
        self.sell_threshold = sell_threshold
        self.start_date = start_date
        self.end_date = end_date
        self.enable_backtest = enable_backtest
    
    def filter_stocks(self, stock_data):
        """
        筛选符合条件的股票
        
        参数:
            stock_data: 包含股票数据的DataFrame
        
        返回:
            符合条件的股票DataFrame
        """
        return stock_data[
            (stock_data['pe'] >= self.pe_threshold) & 
            (stock_data['dividend_yield'] >= self.dividend_yield)
        ]
    
    def calculate_ma120(self, price_data):
        """计算120日移动平均线"""
        return price_data.rolling(window=self.ma_period).mean()
    
    def get_signal(self, price, ma120):
        """
        获取交易信号
        
        返回:
            1: 买入信号
            -1: 卖出信号
            0: 无信号
        """
        if price < ma120 * (1 - self.buy_threshold):
            return 1
        elif price > ma120 * (1 + self.sell_threshold):
            return -1
        return 0
    
    def load_historical_data(self):
        """
        从AKShare加载历史股票数据
        """
        all_stock_data = pd.DataFrame()
        for code in self.stock_codes:
            print(f"正在获取股票 {code} 的历史数据... (start_date: {self.start_date}, end_date: {self.end_date})")
            df = get_stock_zh_a_hist_safe1(
                symbol=code,
                start_date=self.start_date.replace('-', ''),
                end_date=self.end_date.replace('-', ''),
                adjust="qfq"
            )
            if df is not None and not df.empty:
                df['code'] = code
                df['日期'] = pd.to_datetime(df['日期'])
                df.set_index('日期', inplace=True)
                df.rename(columns={'开盘': 'open', '收盘': 'close', '最高': 'high', '最低': 'low', '成交量': 'volume'}, inplace=True)
                all_stock_data = pd.concat([all_stock_data, df[['code', 'open', 'close', 'high', 'low', 'volume']]])
            time.sleep(1) # 避免请求过快
        return all_stock_data

    def load_stock_basics(self):
        """
        模拟加载股票基本面数据 (市盈率和股息率)
        实际应用中需要从数据源获取
        """
        # 这里使用模拟数据，实际应从数据接口获取
        basics = {
            '600900': {'name': '长江电力', 'pe': 25, 'dividend_yield': 0.035},
            '601919': {'name': '中远海控', 'pe': 22, 'dividend_yield': 0.04},
            # 添加更多股票的基本面数据
        }
        return pd.DataFrame.from_dict(basics, orient='index')

    def backtest(self):
        """
        策略回测
        """
        historical_data = self.load_historical_data()
        stock_basics = self.load_stock_basics()

        if historical_data.empty:
            print("未能加载到足够的历史数据，回测中止。")
            return pd.DataFrame()
        if stock_basics.empty:
            print("未能加载到足够的基本面数据，回测中止。")
            return pd.DataFrame()
        print(f"成功加载 {len(historical_data)} 条历史数据和 {len(stock_basics)} 条基本面数据。")

        # 合并历史数据和基本面数据
        # 注意：这里简化处理，假设基本面数据在整个回测期间不变
        # 实际应用中，基本面数据是随时间变化的
        merged_data = historical_data.reset_index().merge(
            stock_basics.reset_index().rename(columns={'index': 'code'}),
            on='code', how='left'
        ).set_index('日期')

        # 筛选符合基本面条件的股票
        filtered_data = self.filter_stocks(merged_data.copy())
        if filtered_data.empty:
            print("没有股票符合基本面筛选条件，回测中止。")
            return pd.DataFrame()
        print(f"经过基本面筛选后，剩余 {len(filtered_data['code'].unique())} 只股票。")

        results = []
        for code in self.stock_codes:
            stock_data = filtered_data[filtered_data['code'] == code].copy()
            if stock_data.empty:
                continue
            stock_data['ma120'] = self.calculate_ma120(stock_data['close'])
            stock_data['signal'] = stock_data.apply(
                lambda x: self.get_signal(x['close'], x['ma120']), axis=1)

            # 计算收益
            stock_data['return'] = stock_data['close'].pct_change().fillna(0)
            stock_data['strategy_return'] = stock_data['return'] * stock_data['signal'].shift(1).fillna(0)

            results.append(stock_data)

        return pd.concat(results) if results else pd.DataFrame()
    
    def generate_report(self, backtest_results, stock_basics):
        """
        生成策略报告
        
        返回:
            报告字符串
        """
        report = "红利白马股策略回测报告\n"
        report += f"回测期间: {self.start_date} 至 {self.end_date}\n"
        report += f"参数设置: PE>={self.pe_threshold}, 股息率>={self.dividend_yield*100}%"
        report += f", MA{self.ma_period}, 买入阈值={self.buy_threshold*100}%, 卖出阈值={self.sell_threshold*100}%\n"
        
        for code in self.stock_codes:
            stock_data = backtest_results[backtest_results['code'] == code]
            stock_name = stock_basics.loc[code, 'name'] if code in stock_basics.index else code
            
            report += f"\n### 股票 {stock_name} ({code}):\n"
            report += f"- 最近买入点: {stock_data[stock_data['signal'] == 1].iloc[-1]['close'] if any(stock_data['signal'] == 1) else '无'}\n"
            report += f"- 最近卖出点: {stock_data[stock_data['signal'] == -1].iloc[-1]['close'] if any(stock_data['signal'] == -1) else '无'}\n"
            
        return report

# 示例用法
if __name__ == "__main__":
    # 可配置参数
    strategy = DividendStrategy(
        stock_codes=ANALYSIS_CONFIG['stock_codes'],
        pe_threshold=ANALYSIS_CONFIG['pe_threshold'],
        dividend_yield=ANALYSIS_CONFIG['dividend_yield'],
        ma_period=ANALYSIS_CONFIG['ma_period'],
        buy_threshold=ANALYSIS_CONFIG['buy_threshold'],
        sell_threshold=ANALYSIS_CONFIG['sell_threshold'],
        start_date=ANALYSIS_CONFIG['start_date'],
        end_date=ANALYSIS_CONFIG['end_date'],
        enable_backtest=ANALYSIS_CONFIG['enable_backtest']
    )
    
    if strategy.enable_backtest:
        print("开始回测...")
        backtest_results = strategy.backtest()

        if not backtest_results.empty:
            stock_basics = strategy.load_stock_basics()
            report = strategy.generate_report(backtest_results, stock_basics)
            print(report)

            # 进一步分析总收益
            total_strategy_return = backtest_results.groupby('code')['strategy_return'].sum()
            print("\n各股票策略总收益:")
            print(total_strategy_return)


            # 输出到带时间的md文档
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            report_filename = f"dividend_strategy_report_{timestamp}.md"
            with open(report_filename, 'w', encoding='utf-8') as f:
                f.write(report)
                f.write("\n\n各股票策略总收益:\n")
                f.write(total_strategy_return.to_string())
                f.write("\n\n---\n优化建议:\n")
                f.write("1. 考虑动态调整PE和股息率阈值，例如根据市场整体情况或行业特点进行调整。\n")
                f.write("2. 引入止损止盈机制，例如当亏损达到一定比例时强制平仓，或盈利达到一定比例时锁定利润。\n")
                f.write("3. 结合更多技术指标，例如RSI、MACD等，增加信号的准确性。\n")
                f.write("4. 考虑交易成本和滑点对收益的影响。\n")
            print(f"回测报告已保存到 {report_filename}")
        else:
            print("回测未能产生结果，请检查数据加载和筛选条件。")
    else:
        print("回测功能已禁用。")
